#include "proxytool.hpp"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cerrno>

int_t Base64EncodeBlock(const char_t* szSrc, char_t* szDst, int_t nLength)
{
	char szTable[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
	switch (nLength)
	{
	case 3:
		{
			szDst[0] = szTable[szSrc[0]>>2];
			szDst[1] = szTable[((szSrc[0]&0x03)<<4) | (szSrc[1]>>4)];
			szDst[2] = szTable[((szSrc[1]&0x0F)<<2) | (szSrc[2]>>6)];
			szDst[3] = szTable[szSrc[2]&0x03F];
		}
		break;
	case 2:
		{
			szDst[0] = szTable[szSrc[0]>>2];
			szDst[1] = szTable[((szSrc[0]&0x03)<<4) | (szSrc[1]>>4)];
			szDst[2] = szTable[(szSrc[1]&0x0F)<<2];
			szDst[3] = '=';
		}
		break;
	case 1:
		{
			szDst[0] = szTable[szSrc[0]>>2];
			szDst[1] = szTable[(szSrc[0]&0x03)<<4];
			szDst[2] = '=';
			szDst[3] = '=';
		}
		break;
	}

	return 0;
}

int_t Base64Encode(const char_t* szUser, const char_t* szPass, char_t* szResult)
{
	char_t szData[1024] = {0};
	sprintf(szData, "%s:%s", szUser, szPass);

	int_t nSrcLength = strlen(szData);
	int_t nTotalLength = nSrcLength%3 ? ((nSrcLength/3+1) * 4) : ((nSrcLength/3)*4);

	memset(szResult, 0, nTotalLength*sizeof(char_t));
	int_t i = 0, j = 0;
	while (i+3 < nSrcLength) {
		Base64EncodeBlock(szData+i, szResult+j, 3);

		i = i + 3;
		j = j + 4;
	}
	Base64EncodeBlock(szData+i, szResult+j, nSrcLength - i);

	return nTotalLength;
}

CProxyTool::CProxyTool ()
{
#ifdef WIN32
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2,2), &wsaData);
#endif
	memset(m_szProxyIp, 0, ADDR_LENGTH*sizeof(char_t));
	memset(m_szUser, 0, NAME_LENGTH*sizeof(char_t));
	memset(m_szPass, 0, PWD_LENGTH*sizeof(char_t));

	m_sProxyPort = -1;
	m_bAuth   = false;
	m_bPass   = false;
}

CProxyTool::~CProxyTool() 
{	
#ifdef WIN32
	WSACleanup();
#endif
}

/****************************************************************************
/* setup proxy server info                    
/****************************************************************************/
int_t CProxyTool::SetProxyServer(const char_t* szProxyIp, ushort_t sProxyPort)
{	
	strncpy(m_szProxyIp, szProxyIp, ADDR_LENGTH);
	m_sProxyPort = sProxyPort;
	m_bAuth = false;
	m_bPass = false;

	return 0;
}
/****************************************************************************
/* setup proxy server auth information   
/****************************************************************************/
int_t CProxyTool::SetProxyAuth(const char_t* szUser, const char_t* szPass)
{	
	strncpy(m_szUser, szUser, NAME_LENGTH);
	strncpy(m_szPass, szPass, PWD_LENGTH);
	m_bAuth = true;
	m_bPass = false;

	return 0;
}

/****************************************************************************
/* connect to remote server
/* this is a public interface for all inherited classed which defined the 
/* steps to connect server. 
/****************************************************************************/
int_t CProxyTool::ConnectServer(const char_t* szServerIp, int_t sServerPort) {
	int_t ret = ConnectProxy();
	if (PT_SUCCESS != ret) {
		DebugPrint("connect proxy server error\n");
		return PT_FAILED;
	}
	ret = Authenticate(szServerIp, sServerPort);
	if (PT_SUCCESS != ret) {
		DebugPrint("authenticate error\n");
		return PT_FAILED;
	}

	return PT_SUCCESS;
}

/****************************************************************************
/* connect to proxy server
/****************************************************************************/
int_t CProxyTool::ConnectProxy() {
	int_t alen = sizeof(struct sockaddr_in);
	struct sockaddr_in proxy_addr;
	
	memset(&proxy_addr, 0, sizeof(proxy_addr));
	proxy_addr.sin_family = AF_INET;
	proxy_addr.sin_addr.s_addr = inet_addr(m_szProxyIp);
	proxy_addr.sin_port = htons(m_sProxyPort);

	m_sSocket = socket(AF_INET, SOCK_STREAM, 0);
	if (m_sSocket < 0) {
		DebugPrint("socket error: %d-%s\n", errno, strerror(errno));
		return PT_SOCKETERR;
	}

	if (connect(m_sSocket, (struct sockaddr*)&proxy_addr, alen) < 0) {
		DebugPrint("connect error: %d-%s\n", errno, strerror(errno));
		return PT_CONNECTERR;
	}

	return PT_SUCCESS;
}

/****************************************************************************
 * disconnect from remote server, default action
 ****************************************************************************/
int_t CProxyTool::DisconnectServer() {
#ifdef WIN32
	m_bPass = false;
	closesocket(m_sSocket);
#else
	m_bPass = false;
	close(m_sSocket);
#endif

	DebugPrint("disconnect from remote server\n");
	return 0;
}

/****************************************************************************
 * disconnect from remote server, default action
 ****************************************************************************/
CHttpProxyTool::CHttpProxyTool(int_t nAuthMethod)
{
	m_nAuthMethod = nAuthMethod;
}

CHttpProxyTool::~CHttpProxyTool() 
{
}

/************************************************************************/
/* no auth method
/************************************************************************/
#define SUCCESS_RESPONSE "200"
int_t CHttpProxyTool::NoAuthMethod(const char_t* szServerIp, int_t sServerPort )
{
	/* 
	 * no auth method used so just build the http header
	 */
	char szAuthInfo[BUFF_LENGTH] = {0};	
	sprintf(szAuthInfo, "CONNECT %s:%d HTTP/1.1\r\n"\
						"Accept: */*\r\n"\
						"Content-Type: text/html\r\n"\
						"Content-Length: 0\r\n"\
						"Proxy-Connection: Keep-Alive\r\n"\
						"User-Agent: Transparenter\r\n\r\n", szServerIp, sServerPort);	
	
	/*
	 * send request to proxy server, if success in auth, it will return http 200 response
	 * so check the return
	 */
	int_t bytes = 0;
	bytes = send(m_sSocket, szAuthInfo, strlen(szAuthInfo), 0);
	if (bytes <= 0) {
		DebugPrint("send error: %d-%s\n", errno, strerror(errno));
		return PT_NETWORKERR;
	}

	memset(szAuthInfo, 0, 1024);
	bytes = recv(m_sSocket, szAuthInfo, BUFF_LENGTH, 0);
	if (bytes <= 0) {
		DebugPrint("recv error: %d-%s\n", errno, strerror(errno));
		return PT_NETWORKERR;
	}

	if (strstr(szAuthInfo, SUCCESS_RESPONSE) != NULL)
		return PT_SUCCESS;
	else
		return PT_FAILED;
}

/************************************************************************/
/* basic auth method for http auth
/* use the base64 encoded username & passwd in the http header
/* connect command http/1.1
/************************************************************************/
int_t CHttpProxyTool::BasicAuthMethod(const char_t* szServerIp, int_t sServerPort)
{	
	if (false == m_bAuth) {
		DebugPrint("not setup auth info\n");
		return PT_NOTSETAUTH;
	}

	/* base64 encode the username & passwd */
	char_t szUserInfo[BUFF_LENGTH] = {0};
	Base64Encode(m_szUser, m_szPass, szUserInfo);

	/* 
	 * build the http header
	 * Proxy-Connection: Keep-Alive
	 * Proxy-Authorization:Basic XXXX
	 */
	char_t szAuthInfo[BUFF_LENGTH] = {0};
	sprintf(szAuthInfo, "CONNECT %s:%d HTTP/1.1\r\n"\
						"Accept: */*\r\n"\
						"Content-Type: text/html\r\n"\
						"Content-Length: 0\r\n"\
						"Proxy-Connection: Keep-Alive\r\n"\
						"Proxy-Authorization:Basic %s\r\n"\
						"User-Agent: Transparenter\r\n\r\n", szServerIp, sServerPort, szUserInfo);
	

	/*
	 * send request to proxy server, if success in auth, it will return http 200 response
	 * so check the return
	 */
	int_t bytes = 0;
	bytes = send(m_sSocket, szAuthInfo, strlen(szAuthInfo), 0);
	if (bytes <= 0) {
		DebugPrint("send error: %d-%s\n", errno, strerror(errno));
		return PT_NETWORKERR;
	}

	memset(szAuthInfo, 0, BUFF_LENGTH);
	bytes = recv(m_sSocket, szAuthInfo, BUFF_LENGTH, 0);
	if (bytes <= 0) {
		DebugPrint("recv error: %d-%s\n", errno, strerror(errno));
		return PT_NETWORKERR;
	}

	if (strstr(szAuthInfo, SUCCESS_RESPONSE) != NULL)
		return PT_SUCCESS;
	else
		return PT_FAILED;
}

/**********************************************************************************
 * template framework
 **********************************************************************************/
int_t CHttpProxyTool::Authenticate(const char_t* szServerIp, int_t sServerPort)
{
	int_t ret = 0;
	switch (m_nAuthMethod)
	{
	case HTTP_NOAUTH:
		ret = NoAuthMethod(szServerIp, sServerPort);
		break;
	case HTTP_BASIC:
		ret = BasicAuthMethod(szServerIp, sServerPort);
		break;
	default:
		ret = PT_NOTSUPPORT;
		break;
	}
	
	if (PT_SUCCESS != ret) {
		DebugPrint("authenticate error: %d\n", ret);
		return PT_FAILED;
	}
	else {
		DebugPrint("authenticate success\n");
		m_bPass = true;
		return PT_SUCCESS;
	}
}

CSocks5ProxyTool::CSocks5ProxyTool()
{
}

CSocks5ProxyTool::~CSocks5ProxyTool()
{
}

int_t CSocks5ProxyTool::AuthenProxy()
{
	/*
	 * Send request 1, client declaring auth methods supported
	 */
	char_t szBuffer[BUFF_LENGTH] = {0};
	szBuffer[0] = SOCKS5_VERSION;
	szBuffer[1] = 2;
	szBuffer[2] = SOCKS5_AUTH_NOAUTH;
	szBuffer[3] = SOCKS5_AUTH_USERPASS;
	send(m_sSocket, szBuffer, 4, 0);

	/*
	 * Recv server's response 1, server choosed one auth method
	 */
	memset(szBuffer, 0, BUFF_LENGTH);
	recv(m_sSocket, szBuffer, BUFF_LENGTH, 0);
	if (szBuffer[0]!=SOCKS5_VERSION) {
		DebugPrint("error request1\n");
		return PT_NETWORKERR;
	}

	// check the response to get auth method type
	if (szBuffer[1] == SOCKS5_AUTH_NOAUTH) {
		DebugPrint("socks5 no auth\n");
		m_bPass = true;
		return PT_SUCCESS;
	}
	else {
		if (m_bAuth == false) {
			DebugPrint("not setup user & passwd\n");
			return PT_NOTSETAUTH;
		}
		else {
			if (szBuffer[1] == SOCKS5_AUTH_USERPASS) {
				memset(szBuffer, 0, 600);
				szBuffer[0] = SOCKS5_AUTH_VER;
				szBuffer[1] = strlen(m_szUser);
				memcpy(&szBuffer[2], m_szUser, strlen(m_szUser));
				szBuffer[2+strlen(m_szUser)] = strlen(m_szPass);
				memcpy(&szBuffer[3+strlen(m_szUser)], m_szPass, strlen(m_szPass));

				send(m_sSocket, szBuffer, 3+strlen(m_szUser)+strlen(m_szPass), 0);
			}
			else {
				DebugPrint("not supported auth method\n");
				return PT_NOTSUPPORT;
			}

			memset(szBuffer, 0, BUFF_LENGTH);
			recv(m_sSocket, szBuffer, BUFF_LENGTH, 0);
			if (szBuffer[0]!=SOCKS5_VERSION || szBuffer[1]!=SOCKS5_RES_SUCCESS)
				return PT_FAILED;
		}
	}

	m_bPass = true;
	return PT_SUCCESS;
}

/************************************************************************/
/* set up channel to communicate
/************************************************************************/
int_t CSocks5ProxyTool::SetupChannel(const char_t* szServerIp, int_t sServerPort)
{	
	/********************************************************************************
	 * send cmd to proxy
	 * CMD_CONNECT: setup tcp channel, just return failure or success
	 * CMD_UDPASSI: setup udp channel, return addr & port on which proxy server bind
	 ********************************************************************************/
	char_t szBuffer[BUFF_LENGTH] = {0};
	szBuffer[0] = SOCKS5_VERSION;
	szBuffer[1] = SOCKS5_CMD_CONNECT;
	szBuffer[2] = SOCKS5_RESERVE;
	szBuffer[3] = SOCKS5_ADDR_IPV4;

	unsigned long  addr = inet_addr(szServerIp);
	unsigned short port = htons(sServerPort);
	memcpy(&szBuffer[4], &addr, 4);
	memcpy(&szBuffer[8], &port, 2);

	send(m_sSocket, szBuffer, 10, 0);
	memset(szBuffer, 0, BUFF_LENGTH);
	recv(m_sSocket, szBuffer, BUFF_LENGTH, 0);

	if (szBuffer[0]!=SOCKS5_VERSION || szBuffer[1]!=SOCKS5_RES_SUCCESS) {
		DebugPrint("setup channel fail\n");
		return PT_FAILED;
	}

	return PT_SUCCESS;
}

/*************************************************************************
 * template design pattern
 *************************************************************************/
int_t CSocks5ProxyTool::Authenticate(const char_t* szServerIp, int_t sServerPort)
{
	int_t ret = 0;

	ret = AuthenProxy();
	if (PT_SUCCESS != ret) {
		DebugPrint("authenticate proxy failed");
		m_bPass = false;
		return PT_FAILED;
	}
	m_bPass = true;

	ret = SetupChannel(szServerIp, sServerPort);
	if (PT_SUCCESS != ret) {
		DebugPrint("setupchannel failed");
		return PT_FAILED;
	}

	return PT_SUCCESS;
}

CProxyTool* CProxyToolFactory::GetProxyTool(int_t type, int_t auth) 
{
	CProxyTool* ptr = NULL;
	switch (type)
	{
	case HTTP_PROXY:
		ptr = new CHttpProxyTool(auth);
		break;
	case SOCKS5_PROXY:
		ptr = new CSocks5ProxyTool();
		break;
	default:
		break;
	}

	return ptr;
}